/////////////////////////////////////////////////
//
//		COPYRIGHT@2008,all right reserved by ax.lyb.lei
//
//		FILE	: BaseFutureMisc.cpp
//
//      PURPOSE : MISC functions set
//
//      AUTHOR  : AX
//
//      DATE    : 2008/10/14
//
////////////////////////////////////////
#include "BaseFutureMisc.h"

#include <atlconv.h>
#include "BaseThread.h"
#ifdef _DEBUG
#include "dblog.h"
#endif
namespace BaseFutureMisc
{
	//version 
	static unsigned int m_dwMajor = 0;
	static unsigned int m_dwMinor = 0;
	unsigned int  VersionMajor()
	{
		return m_dwMajor;
	}
	unsigned int  VersionMinor()
	{
		return m_dwMinor;
	}

	BOOL IsWinNT()
	{
		if(OSVersion() > 0 && OSVersion() <= 3)
			return TRUE;

		return FALSE;
	}

	BOOL IsWin95()
	{
		if(OSVersion() > 3) 
			return TRUE;
		return FALSE;
	}

	int  OSVersion()
	{
		static int nVer = 0;

		if( nVer )
			return nVer;

		OSVERSIONINFO OSversion;

		OSversion.dwOSVersionInfoSize=sizeof(OSVERSIONINFO);

		if (!GetVersionEx(&OSversion))
			return 0;

		m_dwMajor = OSversion.dwMajorVersion;
		m_dwMinor = OSversion.dwMinorVersion;

		switch(OSversion.dwPlatformId)
		{
		case VER_PLATFORM_WIN32_NT:
			if(OSversion.dwMajorVersion==5 && OSversion.dwMinorVersion==0)
			{
				nVer = 1;
				return 1;							    
			}
			if(OSversion.dwMajorVersion==5 && OSversion.dwMinorVersion==1)
			{
				nVer = 1;
				return 1;							    
			}
			if(OSversion.dwMajorVersion==5)
			{
				nVer = 1;
				return 1;							    
			}
			if(OSversion.dwMajorVersion==6)
			{
				nVer = 1;
				return 1;							    
			}
			if(OSversion.dwMajorVersion<=4) 	  
			{
				nVer = 3;
				return 3;							    
			}
		case VER_PLATFORM_WIN32_WINDOWS:
			if(OSversion.dwMinorVersion==0)
			{
				nVer = 5;
				return 5; //95
			}
			nVer = 4;
			return 4;
		}
		return 0;
	}

	//is service?
	static BOOL m_bService = FALSE;
	void SetRunningAsService(BOOL bService)
	{
		m_bService = bService;
	}
	BOOL IsRunningAsService()
	{
		return m_bService;
	}


	BOOL CALLBACK
		KillScreenSaverFunc(HWND hwnd, LPARAM lParam)
	{
		char buffer[256];

		// - ONLY try to close Screen-saver windows!!!
		if ((GetClassNameA(hwnd, buffer, 256) != 0) &&
			(strcmp(buffer, "WindowsScreenSaverClass") == 0))
			PostMessage(hwnd, WM_CLOSE, 0, 0);
		return TRUE;
	}

	void KillScreenSaver()
	{
		if(IsWin95())
		{
			HWND hsswnd = FindWindow (L"WindowsScreenSaverClass", NULL);
			if (hsswnd != NULL)
			{
				PostMessage(hsswnd, WM_CLOSE, 0, 0); 
			}
			return;
		}
		// Windows NT

		// Find the screensaver desktop
		HDESK hDesk = OpenDesktop(
			L"Screen-saver",
			0,
			FALSE,
			DESKTOP_READOBJECTS | DESKTOP_WRITEOBJECTS
			);
		if (hDesk != NULL)
		{ 
			// Close all windows on the screen saver desktop
			EnumDesktopWindows(hDesk, (WNDENUMPROC) &KillScreenSaverFunc, 0);
			CloseDesktop(hDesk);
			// Pause long enough for the screen-saver to close
			//Sleep(2000);
			// Reset the screen saver so it can run again
			SystemParametersInfo(SPI_SETSCREENSAVEACTIVE, TRUE, 0, SPIF_SENDWININICHANGE); 
		}

	}


	BOOL SelectHDESK(void* hDesk)
	{
		// Are we running on NT?
		if (!IsWinNT() || hDesk == NULL)
			return TRUE;

		HDESK hNewDesk = (HDESK)hDesk;
		HDESK hOldDesk = GetThreadDesktop(GetCurrentThreadId());

		DWORD dummy;
		WCHAR new_name[256];

		if (!GetUserObjectInformation(hNewDesk, UOI_NAME, &new_name, 256, &dummy)) {
			return FALSE;
		}

		USES_CONVERSION;

		LOG(LL_INTERR, VNCLOG("SelectHDESK to %s (%x) from %x\n"), W2A(new_name), hNewDesk, hOldDesk);

		// Switch the desktop
		if(!SetThreadDesktop(hNewDesk)) {

			LOG("\n SelectHDESK set thread new desk failed %d.",GetLastError());
			return FALSE;
		}

		// Switched successfully - destroy the old desktop
		if (!CloseDesktop(hOldDesk))
		{
			assert(0);
#ifdef _DEBUG
			LOG("\n close old desk failed %d,error =%d.",hOldDesk,GetLastError());
#endif
		}

		return TRUE;
	}

	BOOL SelectDesktop(const char* szName)
	{
		if(!IsWinNT())
			return (szName == NULL);

		HDESK hDesk = 0;
		WCHAR *name = 0;
		USES_CONVERSION;
		if(szName && strlen(szName) > 0)
		{
			name = A2T(szName);
		}

		if (name != NULL)
		{
			//vnclog.Print(LL_INTERR, VNCLOG("OpenInputdesktop2 named\n"));
			// Attempt to open the named desktop
			hDesk = OpenDesktop(name, 0, FALSE,
				DESKTOP_CREATEMENU | DESKTOP_CREATEWINDOW |
				DESKTOP_ENUMERATE | DESKTOP_HOOKCONTROL |
				DESKTOP_WRITEOBJECTS | DESKTOP_READOBJECTS |
				DESKTOP_SWITCHDESKTOP | GENERIC_WRITE);
		}
		else
		{
			//vnclog.Print(LL_INTERR, VNCLOG("OpenInputdesktop2 NULL\n"));
			// No, so open the input desktop
			hDesk = OpenInputDesktop(0, FALSE,
				DESKTOP_CREATEMENU | DESKTOP_CREATEWINDOW |
				DESKTOP_ENUMERATE | DESKTOP_HOOKCONTROL |
				DESKTOP_WRITEOBJECTS | DESKTOP_READOBJECTS |
				DESKTOP_SWITCHDESKTOP | GENERIC_WRITE);
		}

		// Did we succeed?
		if (hDesk == NULL) {
			//	vnclog.Print(LL_INTERR, VNCLOG("OpenInputdesktop2 \n"));
			assert(0);
			return FALSE;
		}

		LOG("\nOpenInputdesktop2 \n");

		// Switch to the new desktop
		if (!SelectHDESK(hDesk)) 
		{
			// Failed to enter the new desktop, so free it!
			if (!CloseDesktop(hDesk))
			{
				LOG(LL_INTERR, "SelectDesktop failed to close desktop,err = %d.\n",GetLastError());
			}
			return FALSE;
		}

		// We successfully switched desktops!
		return TRUE;
	}

	BOOL GetCurrentUser(char* buf,int nSize)
	{
		//unknown os!
		if( OSVersion () == 0)
		{
			return FALSE;
		}
		if(	IsWinNT && IsRunningAsService())
		{
			// -=- FIRSTLY - verify that a user is logged on

			// Get the current Window station
			HWINSTA hStation = GetProcessWindowStation();
			if (hStation == NULL)
				return FALSE;

			// Get the current user SID size
			DWORD usersize;
			GetUserObjectInformation(hStation,
				UOI_USER_SID, NULL, 0, &usersize);
			DWORD  dwErrorCode = GetLastError();
			SetLastError(0);

			// Check the required buffer size isn't zero
			if (usersize == 0)
			{
				// No user is logged in - ensure we're not impersonating anyone
				RevertToSelf();
				//	g_impersonating_user = FALSE;

				// Return "" as the name ...
				if (strlen("") >= nSize)
					return FALSE;
				strcpy(buf, "");

				return TRUE;
			}

			// -=- SECONDLY - a user is logged on but if we're not impersonating
			//     them then we can't continue!
			//if (!g_impersonating_user) {
			//	// Return "" as the name...
			//	if (strlen("") >= size)
			//		return FALSE;
			//	strcpy(buffer, "");
			//	return TRUE;
			//}

		}

		DWORD length = nSize;

		if (GetUserNameA(buf, &length) == 0)
		{
			UINT error = GetLastError();

			if (error == ERROR_NOT_LOGGED_ON)
			{
				// No user logged on
				if (strlen("") >= nSize)
					return FALSE;
				strcpy(buf, "");
				return TRUE;
			}
			else
			{
				// Genuine error...
				LOG(LL_INTERR, VNCLOG("getusername error %d\n"), GetLastError());
				return FALSE;
			}
		}

		return FALSE;
	}

	// Find the visible window station and switch to it
	// This would allow the service to be started non-interactive
	// Needs more supporting code & a redesign of the server core to
	// work, with better partitioning between server & UI components.

	static HWINSTA home_window_station = GetProcessWindowStation();
	BOOL CALLBACK WinStationEnumProc(LPTSTR name, LPARAM param) {
		HWINSTA station = OpenWindowStation(name, FALSE, GENERIC_ALL);
		HWINSTA oldstation = GetProcessWindowStation();
		USEROBJECTFLAGS flags;
		if (!GetUserObjectInformation(station, UOI_FLAGS, &flags, sizeof(flags), NULL)) {
			return TRUE;
		}
		BOOL visible = flags.dwFlags & WSF_VISIBLE;
		if (visible) {
			if (SetProcessWindowStation(station)) {
				if (oldstation != home_window_station) {
					CloseWindowStation(oldstation);
				}
			} else {
				CloseWindowStation(station);
			}
			return FALSE;
		}
		return TRUE;
	}

	BOOL SelectInputWinStation()
	{
		home_window_station = GetProcessWindowStation();
		return EnumWindowStations(&WinStationEnumProc, NULL);

	}
	void SelectHomeWinStation()
	{
		HWINSTA station=GetProcessWindowStation();
		SetProcessWindowStation(home_window_station);
		CloseWindowStation(station);
	}
	BOOL InputDesktopSelected()
	{
		if (!IsWinNT())
			return TRUE;
		// Get the input and thread desktops
		HDESK threaddesktop = GetThreadDesktop(GetCurrentThreadId());
		HDESK inputdesktop = OpenInputDesktop(0, FALSE,
			DESKTOP_CREATEMENU | DESKTOP_CREATEWINDOW |
			DESKTOP_ENUMERATE | DESKTOP_HOOKCONTROL |
			DESKTOP_WRITEOBJECTS | DESKTOP_READOBJECTS |
			DESKTOP_SWITCHDESKTOP | GENERIC_WRITE);


		// Get the desktop names:
		// *** I think this is horribly inefficient but I'm not sure.
		if (inputdesktop == NULL)
		{
			DWORD lasterror;
			lasterror=GetLastError();
			LOG(LL_INTERR, VNCLOG("OpenInputDesktop I\n"));
			if (lasterror==170) return TRUE;
			LOG(LL_INTERR, VNCLOG("OpenInputDesktop II\n"));
			return FALSE;
		}

		DWORD dummy;
		WCHAR threadname[256];
		WCHAR inputname[256];

		if (!GetUserObjectInformation(threaddesktop, UOI_NAME, &threadname, 256, &dummy))
		{
			if (!CloseDesktop(inputdesktop))
			{
				LOG( VNCLOG("failed to close input desktop\n"));
			}
			return FALSE;
		}
		assert(dummy <= 256);
		if (!GetUserObjectInformation(inputdesktop, UOI_NAME, &inputname, 256, &dummy)) 
		{
			if (!CloseDesktop(inputdesktop))
			{
				LOG( VNCLOG("failed to close input desktop\n"));
			}
			return FALSE;
		}
		assert(dummy <= 256);

		if (!CloseDesktop(inputdesktop))
		{
			LOG(VNCLOG("failed to close input desktop\n"));
		}

		if (wcscmp(threadname, inputname) != 0)
			return FALSE;

		return TRUE;
	}

	// Static routine used to fool Winlogon into thinking CtrlAltDel was pressed
	void *SimulateCtrlAltDelThreadFn(void *context)
	{
		HDESK old_desktop = GetThreadDesktop(GetCurrentThreadId());

		// Switch into the Winlogon desktop
		if (!SelectDesktop("Winlogon"))
		{
			LOG(LL_INTERR, VNCLOG("failed to select logon desktop\n"));
			return NULL;
		}

		LOG(LL_ALL, VNCLOG("generating ctrl-alt-del\n"));

		// Fake a hotkey event to any windows we find there.... :(
		// Winlogon uses hotkeys to trap Ctrl-Alt-Del...
		PostMessage(HWND_BROADCAST, WM_HOTKEY, 0, MAKELONG(MOD_ALT | MOD_CONTROL, VK_DELETE));

		// Switch back to our original desktop
		if (old_desktop != NULL)
			SelectHDESK(old_desktop);

		return NULL;
	}

	class SimulateCtrlAltDelObj
	{
	public:
		int Run(void* lparam)
		{
			SimulateCtrlAltDelThreadFn(lparam);
			return 0;
		}
		void Stop()
		{
		}
	};
	BOOL SimulateCtrlAltDel()
	{
		if (!IsWinNT())
			return TRUE;
		static SimulateCtrlAltDelObj cObj;
		static BaseThread<SimulateCtrlAltDelObj> cThread(cObj);
		cThread.Stop();
		cThread.Start(NULL);
		return TRUE;
	}

	BOOL LockWorkstation()
	{
		if (!IsWinNT()) 
		{
			return FALSE;
		}
		HMODULE user32 = LoadLibrary(L"user32.dll");
		if (!user32) 
		{
			return FALSE;
		}
		typedef BOOL (*LWProc) ();
		LWProc lockworkstation = (LWProc)GetProcAddress(user32, "LockWorkStation");
		if (!lockworkstation) 
		{
			FreeLibrary(user32);
			return FALSE;
		}
		// Attempt to lock the workstation
		BOOL result = (lockworkstation)();

		FreeLibrary(user32);
		if (!result) 
		{
			assert(0);
			return FALSE;
		}
		return TRUE;
	}


	class InstObj
	{
	public:
		InstObj()
		{
			IsWinNT();
			IsWin95();
			OSVersion();
		}
	};

	static InstObj obj;


	//DynamicLib
	DynamicLib::DynamicLib(const char* pszLibName)
	{
		m_hMod = LoadLibraryA(pszLibName);
	}

	DynamicLib::~DynamicLib()
	{
		if(m_hMod)
			FreeLibrary(m_hMod);
	}

	void* DynamicLib::GetProcAddress(const char* pszFunName)
	{
		if(m_hMod == NULL || pszFunName == NULL)
			return NULL;
		return ::GetProcAddress(m_hMod,pszFunName);
	}

	void SessionFix() 
	{

		if (!IsRunningAsService()) return;

		DWORD  dwSessionId;
		DWORD  dwCurrentSessionId;
		typedef DWORD (WINAPI* pWTSGetActiveConsoleSessionId)(VOID);
		typedef BOOL (WINAPI * pProcessIdToSessionId) (DWORD, DWORD*);
		typedef BOOLEAN (WINAPI * pWinStationConnect) (HANDLE,ULONG,ULONG,PCWSTR,ULONG);
		typedef BOOL (WINAPI * pLockWorkStation)();
		HMODULE  hlibkernel = LoadLibrary(L"kernel32.dll"); 
		HMODULE  hlibwinsta = LoadLibrary(L"winsta.dll"); 
		HMODULE  hlibuser32 = LoadLibrary(L"user32.dll"); 
		pWTSGetActiveConsoleSessionId WTSGetActiveConsoleSessionIdF=NULL;
		pProcessIdToSessionId ProcessIdToSessionIdF=NULL;
		pWinStationConnect WinStationConnectF=NULL;
		pLockWorkStation LockWorkStationF=NULL;

		if (hlibkernel) 
		{
			WTSGetActiveConsoleSessionIdF = (pWTSGetActiveConsoleSessionId)\
				GetProcAddress(hlibkernel, "WTSGetActiveConsoleSessionId"); 
			ProcessIdToSessionIdF=(pProcessIdToSessionId)GetProcAddress(hlibkernel, "ProcessIdToSessionId"); 
		}
		if (hlibwinsta)
		{
			WinStationConnectF=(pWinStationConnect)GetProcAddress(hlibwinsta, "WinStationConnectW"); 
		}
		if (hlibuser32)
		{
			LockWorkStationF=(pLockWorkStation)GetProcAddress(hlibuser32, "LockWorkStation"); 
		}

		if (WTSGetActiveConsoleSessionIdF!=NULL)
			dwSessionId =WTSGetActiveConsoleSessionIdF();
		if (ProcessIdToSessionIdF!=NULL)
			ProcessIdToSessionIdF(GetCurrentProcessId(), &dwCurrentSessionId);

		if (WTSGetActiveConsoleSessionIdF!=NULL) if (dwSessionId!=dwCurrentSessionId)
		{
			//Black screen only session 0 can be handled from a service
			if (WinStationConnectF!=NULL) 
			{
				WinStationConnectF(0, dwCurrentSessionId, dwSessionId, L"", 0);
				if (LockWorkStationF!=NULL) LockWorkStationF();
			}
		}

		if (hlibkernel) FreeLibrary(hlibkernel);
		if (hlibwinsta) FreeLibrary(hlibwinsta);
		if (hlibuser32) FreeLibrary(hlibuser32);
	}

};
